extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{
    Data, DeriveInput, Fields, Ident, LitStr, Token,
    parse::{Parse, ParseStream},
    parse_macro_input,
};

// 1. 定义一个结构体来表示我们期望从属性中解析出的参数
struct DeftokenArgs {
    typename: LitStr,
    typevalue: LitStr,
}

// 2. 为这个结构体实现 `Parse` trait，告诉 syn 如何解析它
//    它将负责解析 `typename = "...", typevalue = "..."` 这样的内容
impl Parse for DeftokenArgs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut typename = None;
        let mut typevalue = None;

        while !input.is_empty() {
            let key: Ident = input.parse()?;
            input.parse::<Token![=]>()?;
            let value: LitStr = input.parse()?;

            if key == "typename" {
                typename = Some(value);
            } else if key == "typevalue" {
                typevalue = Some(value);
            } else {
                return Err(syn::Error::new(
                    key.span(),
                    "Unknown attribute key. Expected `typename` or `typevalue`.",
                ));
            }

            if input.peek(Token![,]) {
                input.parse::<Token![,]>()?;
            }
        }

        let typename =
            typename.ok_or_else(|| input.error("Missing required attribute key `typename`"))?;
        let typevalue =
            typevalue.ok_or_else(|| input.error("Missing required attribute key `typevalue`"))?;

        Ok(DeftokenArgs {
            typename,
            typevalue,
        })
    }
}

#[proc_macro_derive(SysYToken, attributes(deftoken))]
pub fn deftoken_derive(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    let name = &ast.ident;
    let trait_name = Ident::new(&format!("{name}Info"), name.span());

    // 3. 简化变体数据的提取过程
    let variants = match &ast.data {
        Data::Enum(data) => &data.variants,
        _ => panic!("#[derive(SysYToken)] can only be used on enums"),
    };

    let mut typename_arms = Vec::new();
    let mut typevalue_arms = Vec::new();
    // 为 from_str 实现收集 match arms
    let mut from_str_arms = Vec::new();

    for variant in variants {
        let variant_ident = &variant.ident;
        let attr = variant
            .attrs
            .iter()
            .find(|a| a.path().is_ident("deftoken"))
            .expect("Each variant must have a #[deftoken] attribute");

        // 4. 使用我们定义的解析器来解析属性
        let args: DeftokenArgs = match attr.parse_args() {
            Ok(args) => args,
            Err(e) => return e.to_compile_error().into(),
        };
        let typename = args.typename;
        let typevalue = args.typevalue;

        // --- 生成 Trait 的 match arms ---

        // Typename arm
        let pattern_any = match &variant.fields {
            Fields::Unit => quote! { Self::#variant_ident },
            Fields::Unnamed(..) => quote! { Self::#variant_ident(..) },
            Fields::Named(..) => quote! { Self::#variant_ident { .. } },
        };
        typename_arms.push(quote! { #pattern_any => #typename });

        // Typevalue arm
        match &variant.fields {
            Fields::Unit => {
                typevalue_arms.push(quote! { Self::#variant_ident => #typevalue.to_string() });

                // NEW: 为 unit 变体生成 from_str arm
                // 只有当变体没有字段时，我们才能从一个固定的字符串转换而来
                from_str_arms.push(quote! { #typevalue => Ok(Self::#variant_ident) });
            }
            Fields::Unnamed(fields) => {
                let bindings: Vec<_> = (0..fields.unnamed.len())
                    .map(|i| Ident::new(&format!("v{i}"), variant_ident.span()))
                    .collect();

                let pattern_bind = quote! { Self::#variant_ident(#( #bindings ),*) };

                let mut replace_logic = quote! { #typevalue.to_string() };
                for (i, binding) in bindings.iter().enumerate() {
                    let placeholder = format!("${i}");
                    replace_logic =
                        quote! { #replace_logic.replace(#placeholder, &#binding.to_string()) };
                }

                typevalue_arms.push(quote! { #pattern_bind => { #replace_logic } });
            }
            Fields::Named(..) => panic!("Named fields are not supported in this macro version."),
        }
    }

    // 组合最终生成的代码
    let trait_impl = quote! {
        pub trait #trait_name {
            fn get_typename(&self) -> &'static str;
            fn get_tokenvalue(&self) -> String;
        }

        impl #trait_name for #name {
            fn get_typename(&self) -> &'static str {
                match self {
                    #( #typename_arms ),*
                }
            }

            fn get_tokenvalue(&self) -> String {
                match self {
                    #( #typevalue_arms ),*
                }
            }
        }
    };

    // 生成 FromStr 实现
    let from_str_impl = quote! {
        #[derive(Debug, PartialEq, Eq)]
        pub struct ParseTokenError;

        impl std::fmt::Display for ParseTokenError {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "Provided string does not match any token variant")
            }
        }

        impl std::error::Error for ParseTokenError {}

        impl std::str::FromStr for #name {
            type Err = ParseTokenError;

            fn from_str(s: &str) -> Result<Self, Self::Err> {
                match s {
                    #( #from_str_arms ),*,
                    // 如果没有匹配的变体，返回错误
                    // 注意：这里我们只能匹配那些没有字段的变体
                    _ => Err(ParseTokenError),
                }
            }
        }
    };

    // 合并所有生成的代码
    let final_expanded = quote! {
        #trait_impl
        #from_str_impl
    };

    TokenStream::from(final_expanded)
}
